# dynamic_noise_optimizer.py
import numpy as np

class DynamicNoiseOptimizer:
    """
    Dynamically adjusts noise parameters based on system state
    using federated clustering and reinforcement learning (simplified)
    """
    
    def __init__(self, min_sigma=0.5, max_sigma=2.0):
        self.min_sigma = min_sigma
        self.max_sigma = max_sigma
        self.cluster_threshold = 0.8
    
    def adapt_noise_parameters(self, gradient_variance, bandwidth, energy_level):
        """
        Adapt noise parameters based on system state
        
        Args:
            gradient_variance (float): Measure of data heterogeneity
            bandwidth (float): Current network bandwidth (Mbps)
            energy_level (float): Client device energy status
            
        Returns:
            float: Optimized noise sigma value
        """
        if gradient_variance > 0.5 or energy_level < 0.3:
            return self.min_sigma  # Reduce noise for high heterogeneity/low energy
        elif bandwidth < 50:
            return self.min_sigma  # Reduce noise for low bandwidth
        else:
            return self.max_sigma  # Use stronger noise otherwise
    
    def update_clustering_threshold(self, round_num, total_rounds):
        """
        Dynamically adjust federated clustering threshold
        
        Args:
            round_num (int): Current training round
            total_rounds (int): Total number of training rounds
        """
        self.cluster_threshold = 0.8 + 0.1 * np.tanh(round_num / total_rounds)
